#!/usr/bin/env python3
"""
设备配置切换脚本
用于在CPU和GPU(MPS)模式之间快速切换FunASR配置
"""

import argparse
import re
from pathlib import Path


class DeviceConfigSwitcher:
    def __init__(self):
        self.config_file = Path("src/speech_convert/config/funasr_config.py")
        self.backup_file = Path("src/speech_convert/config/funasr_config.py.backup")

    def backup_config(self):
        """备份当前配置文件"""
        if self.config_file.exists():
            import shutil

            shutil.copy2(self.config_file, self.backup_file)
            print(f"✅ 配置文件已备份到: {self.backup_file}")
        else:
            print(f"❌ 配置文件不存在: {self.config_file}")
            return False
        return True

    def restore_config(self):
        """恢复配置文件"""
        if self.backup_file.exists():
            import shutil

            shutil.copy2(self.backup_file, self.config_file)
            print(f"✅ 配置文件已从备份恢复")
        else:
            print(f"❌ 备份文件不存在: {self.backup_file}")
            return False
        return True

    def read_config(self):
        """读取配置文件内容"""
        try:
            with open(self.config_file, "r", encoding="utf-8") as f:
                return f.read()
        except Exception as e:
            print(f"❌ 读取配置文件失败: {e}")
            return None

    def write_config(self, content: str):
        """写入配置文件内容"""
        try:
            with open(self.config_file, "w", encoding="utf-8") as f:
                f.write(content)
            return True
        except Exception as e:
            print(f"❌ 写入配置文件失败: {e}")
            return False

    def get_current_device(self):
        """获取当前设备配置"""
        content = self.read_config()
        if not content:
            return None

        # 查找DEFAULT_DEVICE设置 - 支持多行定义
        device_match = re.search(
            r'DEFAULT_DEVICE\s*=\s*\(\s*["\']([^"\']+)["\']',
            content,
            re.MULTILINE | re.DOTALL,
        )
        if not device_match:
            # 尝试单行定义
            device_match = re.search(
                r'DEFAULT_DEVICE\s*=\s*["\']([^"\']+)["\']', content
            )

        ngpu_match = re.search(r"DEFAULT_NGPU\s*=\s*(\d+)", content)

        if device_match:
            device = device_match.group(1)
            ngpu = int(ngpu_match.group(1)) if ngpu_match else 1
            return {
                "device": device,
                "ngpu": ngpu,
                "mode": "GPU" if device in ["mps", "cuda"] and ngpu > 0 else "CPU",
            }
        return None

    def switch_to_cpu(self):
        """切换到CPU模式"""
        content = self.read_config()
        if not content:
            return False

        print("🔄 正在切换到CPU模式...")

        # 替换设备配置 - 处理多行格式
        content = re.sub(
            r'DEFAULT_DEVICE\s*=\s*\(\s*["\'][^"\']+["\']\s*[^)]*\)',
            'DEFAULT_DEVICE = "cpu"',
            content,
            flags=re.MULTILINE | re.DOTALL,
        )
        # 如果是单行格式
        content = re.sub(
            r'DEFAULT_DEVICE\s*=\s*["\'][^"\']+["\']', 'DEFAULT_DEVICE = "cpu"', content
        )
        content = re.sub(r"DEFAULT_NGPU\s*=\s*\d+", "DEFAULT_NGPU = 0", content)

        if self.write_config(content):
            print("✅ 已切换到CPU模式")
            print("   - DEFAULT_DEVICE = 'cpu'")
            print("   - DEFAULT_NGPU = 0")
            return True
        return False

    def switch_to_gpu(self):
        """切换到GPU模式"""
        content = self.read_config()
        if not content:
            return False

        print("🔄 正在切换到GPU模式...")

        # 替换设备配置 - 处理多行格式
        content = re.sub(
            r'DEFAULT_DEVICE\s*=\s*\(\s*["\'][^"\']+["\']\s*[^)]*\)',
            'DEFAULT_DEVICE = "mps"',
            content,
            flags=re.MULTILINE | re.DOTALL,
        )
        # 如果是单行格式
        content = re.sub(
            r'DEFAULT_DEVICE\s*=\s*["\'][^"\']+["\']', 'DEFAULT_DEVICE = "mps"', content
        )
        content = re.sub(r"DEFAULT_NGPU\s*=\s*\d+", "DEFAULT_NGPU = 1", content)

        if self.write_config(content):
            print("✅ 已切换到GPU模式")
            print("   - DEFAULT_DEVICE = 'mps'")
            print("   - DEFAULT_NGPU = 1")
            return True
        return False

    def show_current_config(self):
        """显示当前配置"""
        config = self.get_current_device()
        if config:
            print(f"📊 当前配置:")
            print(f"   - 模式: {config['mode']}")
            print(f"   - 设备: {config['device']}")
            print(f"   - GPU数量: {config['ngpu']}")
        else:
            print("❌ 无法读取当前配置")

    def check_device_support(self):
        """检查设备支持情况"""
        print("🔍 检查设备支持情况...")

        try:
            import torch

            print(f"   - PyTorch版本: {torch.__version__}")
            print(f"   - CUDA可用: {'✅' if torch.cuda.is_available() else '❌'}")
            print(
                f"   - MPS可用: {'✅' if torch.backends.mps.is_available() else '❌'}"
            )
            print(f"   - MPS已构建: {'✅' if torch.backends.mps.is_built() else '❌'}")

            if torch.backends.mps.is_available():
                print("   - 建议使用MPS(GPU)模式以获得更好性能")
            else:
                print("   - 建议使用CPU模式")

        except ImportError:
            print("   - ❌ PyTorch未安装，无法检查GPU支持")


def main():
    parser = argparse.ArgumentParser(description="FunASR设备配置切换工具")
    parser.add_argument(
        "--mode",
        choices=["cpu", "gpu", "status", "check", "backup", "restore"],
        help="操作模式",
    )
    parser.add_argument("--auto", action="store_true", help="自动选择最佳设备")

    args = parser.parse_args()

    switcher = DeviceConfigSwitcher()

    # 如果没有指定参数，显示帮助信息
    if not args.mode and not args.auto:
        print("🎯 FunASR设备配置切换工具")
        print("=" * 50)
        switcher.show_current_config()
        print("\n📖 使用方法:")
        print("   python switch_device_config.py --mode cpu     # 切换到CPU模式")
        print("   python switch_device_config.py --mode gpu     # 切换到GPU模式")
        print("   python switch_device_config.py --mode status  # 显示当前配置")
        print("   python switch_device_config.py --mode check   # 检查设备支持")
        print("   python switch_device_config.py --mode backup  # 备份配置文件")
        print("   python switch_device_config.py --mode restore # 恢复配置文件")
        print("   python switch_device_config.py --auto         # 自动选择最佳设备")
        return

    # 处理自动模式
    if args.auto:
        print("🤖 自动选择最佳设备配置...")
        try:
            import torch

            if torch.backends.mps.is_available():
                print("✅ 检测到MPS支持，切换到GPU模式")
                switcher.switch_to_gpu()
            else:
                print("⚠️ 未检测到MPS支持，切换到CPU模式")
                switcher.switch_to_cpu()
        except ImportError:
            print("⚠️ PyTorch未安装，切换到CPU模式")
            switcher.switch_to_cpu()
        return

    # 处理具体操作
    if args.mode == "cpu":
        if not switcher.backup_config():
            print("❌ 备份失败，操作中止")
            return
        switcher.switch_to_cpu()

    elif args.mode == "gpu":
        if not switcher.backup_config():
            print("❌ 备份失败，操作中止")
            return
        switcher.switch_to_gpu()

    elif args.mode == "status":
        switcher.show_current_config()

    elif args.mode == "check":
        switcher.check_device_support()

    elif args.mode == "backup":
        switcher.backup_config()

    elif args.mode == "restore":
        switcher.restore_config()

    print("\n🔄 操作完成！重启服务器以使配置生效。")


if __name__ == "__main__":
    main()
